package frog.lazy.framework.datasource;

import cn.hutool.core.util.StrUtil;
import cn.hutool.log.StaticLog;
import frog.lazy.framework.core.tools.StrUtils;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.MutablePropertyValues;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.beans.factory.support.GenericBeanDefinition;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.ImportBeanDefinitionRegistrar;
import org.springframework.core.annotation.Order;
import org.springframework.core.env.Environment;
import org.springframework.core.type.AnnotationMetadata;

import javax.annotation.PreDestroy;
import javax.sql.DataSource;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.Executors;

/**
 * @author lwy
 */
@Slf4j
public class DynamicDataSourceRegister implements ImportBeanDefinitionRegistrar , EnvironmentAware {


    /**
     * 指定默认数据源(springboot2.0默认数据源是hikari如何想使用其他数据源可以自己配置)
     */
    private static final String DATASOURCE_TYPE_DEFAULT = "com.zaxxer.hikari.HikariDataSource";
    /**
     * 默认数据源
     */
    private DataSource defaultDataSource;
    /**
     * 用户自定义数据源
     */
    private Map<String, DataSource> slaveDataSources = new HashMap<>();

    @Override
    public void registerBeanDefinitions(AnnotationMetadata annotationMetadata, BeanDefinitionRegistry beanDefinitionRegistry) {
        Map<Object, Object> targetDataSources = new HashMap<Object, Object>();
        //添加默认数据源
        targetDataSources.put("primary", this.defaultDataSource);
        DynamicDataSourceContextHolder.dataSourceIds.add("primary");
        //添加其他数据源
        targetDataSources.putAll(slaveDataSources);
        for (String key : slaveDataSources.keySet()) {
            DynamicDataSourceContextHolder.dataSourceIds.add(key);
        }

        //创建DynamicDataSource
        GenericBeanDefinition beanDefinition = new GenericBeanDefinition();
        beanDefinition.setBeanClass(DynamicDataSource.class);
        beanDefinition.setSynthetic(true);
        MutablePropertyValues mpv = beanDefinition.getPropertyValues();
        mpv.addPropertyValue("defaultTargetDataSource", defaultDataSource);
        mpv.addPropertyValue("targetDataSources", targetDataSources);
        //注册 - BeanDefinitionRegistry
        beanDefinitionRegistry.registerBeanDefinition("dataSource", beanDefinition);

    }

    public DataSource buildDataSource(Map<String, Object> dataSourceMap) {
        try {
            Object type = dataSourceMap.get("type");
            if (type == null) {
                type = DATASOURCE_TYPE_DEFAULT;
            }
            Class<? extends DataSource> dataSourceType;
            dataSourceType = (Class<? extends DataSource>) Class.forName((String) type);
            String driverClassName = dataSourceMap.get("driver").toString();
            String url = dataSourceMap.get("url").toString();
            String username = dataSourceMap.get("username").toString();
            String password = dataSourceMap.get("password").toString();
            boolean readOnly = StrUtils.isBlank(dataSourceMap.get("readOnly")+"") ? true:Boolean.valueOf((String)dataSourceMap.get("readOnly"));
            // 自定义DataSource配置
            DataSourceBuilder factory = DataSourceBuilder.create()
                    .driverClassName(driverClassName)
                    .url(url)
                    .username(username)
                    .password(password)
                    .type(dataSourceType);
            DataSource dataSource = factory.build();
            dataSource.getConnection().setNetworkTimeout(Executors.newFixedThreadPool(10), 2);
            dataSource.getConnection().setReadOnly(readOnly);
            return dataSource;
        } catch (ClassNotFoundException e) {
            e.printStackTrace();
        } catch (SQLException e) {
            e.printStackTrace();
        }
        return null;
    }

    @Override
    public void setEnvironment(Environment environment) {
        StaticLog.info("根据配置读取链接进行初始化");
        initDefaultDataSource(environment);
        initslaveDataSources(environment);
    }

    /**
     * 主库 创建
     * @param env
     */
    private void initDefaultDataSource(Environment env) {
        // 读取主数据源
        Map<String, Object> dsMap = new HashMap<>();
        dsMap.put("driver", env.getProperty("hikari.driver-class-name"));
        dsMap.put("url", env.getProperty("hikari.primary.url"));
        dsMap.put("username", env.getProperty("hikari.primary.username"));
        dsMap.put("password", env.getProperty("hikari.primary.password"));
        dsMap.put("readOnly",env.getProperty("hikari.primary.read-only"));
        defaultDataSource = buildDataSource(dsMap);
    }


    /**
     * 从只读库 创建
     * @param env
     */
    private void initslaveDataSources(Environment env) {
        // 读取配置文件获取更多数据源
        String dsPrefixs = env.getProperty("hikari.slave-names");
        String driver =   env.getProperty("hikari.driver-class-name");
        if(StrUtil.isNotBlank(dsPrefixs)) {
            for (String dsPrefix : dsPrefixs.split(",")) {
                // 多个数据源
                Map<String, Object> dsMap = new HashMap<>();
                dsMap.put("driver", driver);
                dsMap.put("url", env.getProperty("hikari." + dsPrefix + ".url"));
                dsMap.put("username", env.getProperty("hikari." + dsPrefix + ".username"));
                dsMap.put("password", env.getProperty("hikari." + dsPrefix + ".password"));
                dsMap.put("readOnly",env.getProperty("hikari." + dsPrefix + ".read-only"));
                DataSource ds = buildDataSource(dsMap);
                slaveDataSources.put(dsPrefix, ds);
            }
        }
    }

    @PreDestroy
    @Order(99)
    @SneakyThrows
    public void closeConnection(){
        log.debug("关闭数据库链接开始");
        defaultDataSource.getConnection().close();
        Iterator<Map.Entry<String, DataSource>> iterator = slaveDataSources.entrySet().iterator();
        while (iterator.hasNext()){
            iterator.next().getValue().getConnection().close();
        }
        log.debug("关闭数据库链接结束");
    }
}
